#include "eigen/conv2d.h"
#include "eigen/normal.h"
#include "eigen/padding.h"
#include "eigen/conv.h"

namespace  ldl_eigen
{
Conv2D::Conv2D(int64_t in_channels, int64_t out_channels, int64_t img_width, 
    int64_t img_height, int64_t kernel_size, int64_t stride, int64_t padding)
{
    m_in_channels = in_channels;
    m_out_channels = out_channels;
    m_img_width = img_width;
    m_img_height = img_height;
    m_kernel_size = kernel_size;
    m_stride = stride;
    m_padding = padding;

    m_kernels.resize(in_channels * out_channels);
    for(auto &item : m_kernels)
    {
        item = Normal::matrix(m_kernel_size, m_kernel_size, 0.0f, 0.01);
    }

    m_kernel_gradients.resize(in_channels * out_channels);
    for(auto &item : m_kernel_gradients)
    {
        item.resize(m_kernel_size, m_kernel_size);
    }


    m_bias.resize(in_channels * out_channels);
    for(auto &item : m_bias)
    {
        item = 0.0f;
    }

    m_bias_gradient.resize(in_channels * out_channels);

    m_padding_img_width = m_img_width + 2 * m_padding;
    m_padding_img_height = m_img_height + 2*padding;
    m_row_sliding_max = m_padding_img_height - m_kernels.at(0).rows() + 1;
    m_col_sliding_max = m_padding_img_width - m_kernels.at(0).cols() + 1;
    m_output_height = (m_row_sliding_max - 1) / m_stride + 1;
    m_output_width = (m_col_sliding_max - 1) / m_stride + 1;
}

void Conv2D::forward()
{
    // 一个输入通道会生成多个输出通道
    m_output.resize(m_ptr_input->rows() * m_out_channels, m_output_height * m_output_width);

    auto batch_size = m_ptr_input->rows() / m_in_channels;
    for(int64_t batch_index = 0; batch_index < batch_size; batch_index++)
    {
        for(int64_t in_channel = 0;in_channel < m_in_channels;in_channel++)
        {
            auto& single_intput = m_ptr_input->block(batch_index*m_in_channels+in_channel, 0, 1, m_ptr_input->cols());
            for(int64_t out_channel = 0;out_channel<m_out_channels;out_channel++)
            {
                // Conv::conv(single_intput,
                //     m_kernels.at(out_channel*m_in_channels+in_channel), m_bias.at(out_channel*m_in_channels+in_channel), m_stride, 
                //     m_output.block(batch_index*m_in_channels*m_out_channels, 0, 1, m_output.cols()));
            }
        }
    }
}

void Conv2D::backward()
{
    // TODO
}

void Conv2D::update()
{
    // TODO
}

void Conv2D::set_bias(const std::vector<float> bias)
{
    m_bias = bias;
}

void Conv2D::set_kernels(const std::vector<Eigen::MatrixXf> &kernels)
{
    m_kernels = kernels;
}
}